// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/exceptions.h>
#include <pollux/vector/builder_type_utils.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    template<typename T>
    void DictionaryVector<T>::setInternalState() {
        rawIndices_ = indices_->as<vector_size_t>();

        // Sanity check indices for non-null positions. Enabled in debug mode only to
        // avoid performance hit in production.
#ifndef NDEBUG
        validate({});
#endif

        if (is_lazy_not_loaded(*dictionaryValues_)) {
            POLLUX_CHECK(
                dictionaryValues_->mark_as_containing_lazy_and_wrapped(),
                "An unloaded lazy vector cannot be wrapped by two different"
                " top level vectors.");
            // Do not load Lazy vector
            return;
        }
        // Ensure any internal state in dictionaryValues_ is initialized, and it
        // points to the loaded vector underneath any lazy layers.
        dictionaryValues_ = BaseVector::loaded_vector_shared(dictionaryValues_);

        if (dictionaryValues_->is_scalar()) {
            scalarDictionaryValues_ =
                    reinterpret_cast<SimpleVector<T> *>(dictionaryValues_->loaded_vector());
            if (scalarDictionaryValues_->is_flat_encoding() && !std::is_same_v<T, bool>) {
                rawDictionaryValues_ =
                        reinterpret_cast<FlatVector<T> *>(scalarDictionaryValues_)
                        ->rawValues();
            }
        }
        initialized_ = true;

        BaseVector::in_memory_bytes_ =
                BaseVector::nulls_ ? BaseVector::nulls_->capacity() : 0;
        BaseVector::in_memory_bytes_ += indices_->capacity();
        BaseVector::in_memory_bytes_ += dictionaryValues_->in_memory_bytes();
    }

    template<typename T>
    DictionaryVector<T>::DictionaryVector(
        pollux::memory::MemoryPool *pool,
        BufferPtr nulls,
        size_t length,
        VectorPtr dictionaryValues,
        BufferPtr dictionaryIndices,
        const SimpleVectorStats<T> &stats,
        std::optional<vector_size_t> distinctValueCount,
        std::optional<vector_size_t> nullCount,
        std::optional<bool> isSorted,
        std::optional<ByteCount> represented_bytes,
        std::optional<ByteCount> storageByteCount)
        : SimpleVector<T>(
            pool,
            dictionaryValues->type(),
            VectorEncoding::Simple::DICTIONARY,
            nulls,
            length,
            stats,
            distinctValueCount,
            nullCount,
            isSorted,
            represented_bytes,
            storageByteCount) {
        POLLUX_CHECK(dictionaryValues != nullptr, "dictionaryValues must not be null");
        POLLUX_CHECK(
            dictionaryIndices != nullptr, "dictionaryIndices must not be null");
        POLLUX_CHECK_GE(
            dictionaryIndices->size(),
            length * sizeof(vector_size_t),
            "Malformed dictionary, index array is shorter than DictionaryVector");
        dictionaryValues_ = dictionaryValues;
        indices_ = dictionaryIndices;
        setInternalState();
    }

    template<typename T>
    bool DictionaryVector<T>::is_null_at(vector_size_t idx) const {
        POLLUX_DCHECK(initialized_);
        if (BaseVector::is_null_at(idx)) {
            return true;
        }
        auto innerIndex = getDictionaryIndex(idx);
        return dictionaryValues_->is_null_at(innerIndex);
    }

    template<typename T>
    const T DictionaryVector<T>::valueAtFast(vector_size_t idx) const {
        POLLUX_DCHECK(initialized_);
        if (rawDictionaryValues_) {
            return rawDictionaryValues_[getDictionaryIndex(idx)];
        }
        return scalarDictionaryValues_->value_at(getDictionaryIndex(idx));
    }

    template<>
    inline const bool DictionaryVector<bool>::valueAtFast(vector_size_t idx) const {
        POLLUX_DCHECK(initialized_);
        auto innerIndex = getDictionaryIndex(idx);
        return scalarDictionaryValues_->value_at(innerIndex);
    }

    template<typename T>
    std::unique_ptr<SimpleVector<uint64_t> > DictionaryVector<T>::hash_all() const {
        // TODO (T58177479) - optimize to reuse the index vector allowing us to only
        // worry about hashing the dictionary values themselves.  Challenge is that
        // the null information is a part of the index array and so does not have a
        // "null" representation in the dictionary.
        // TODO T70734527 dealing with zero length vector
        if (BaseVector::length_ == 0) {
            return nullptr;
        }
        // If there is at least one value, then indices_ is set and has a pool.
        BufferPtr hashes =
                AlignedBuffer::allocate<uint64_t>(BaseVector::length_, indices_->pool());
        uint64_t *rawHashes = hashes->asMutable<uint64_t>();
        for (vector_size_t i = 0; i < BaseVector::length_; ++i) {
            if (BaseVector::is_null_at(i)) {
                rawHashes[i] = BaseVector::kNullHash;
            } else {
                rawHashes[i] = dictionaryValues_->hash_value_at(rawIndices_[i]);
            }
        }
        return std::make_unique<FlatVector<uint64_t> >(
            BaseVector::pool_,
            BIGINT(),
            BufferPtr(nullptr),
            BaseVector::length_,
            hashes,
            std::vector<BufferPtr>(0) /* stringBuffers */,
            SimpleVectorStats<uint64_t>{},
            BaseVector::distinctValueCount_.value() +
            (BaseVector::nullCount_.value_or(0) > 0 ? 1 : 0),
            0 /* nullCount */,
            false /* sorted */,
            sizeof(uint64_t) * BaseVector::length_ /* represented_bytes */);
    }

    template<typename T>
    ksimd::batch<T> DictionaryVector<T>::loadSIMDValueBufferAt(
        size_t byteOffset) const {
        if constexpr (can_simd) {
            constexpr int N = ksimd::batch<T>::size;
            alignas(ksimd::default_arch::alignment()) T tmp[N];
            auto startIndex = byteOffset / sizeof(T);
            for (int i = 0; i < N; ++i) {
                tmp[i] = valueAtFast(startIndex + i);
            }
            return ksimd::load_aligned(tmp);
        } else {
            POLLUX_UNREACHABLE();
        }
    }

    template<typename T>
    VectorPtr DictionaryVector<T>::slice(vector_size_t offset, vector_size_t length)
    const {
        POLLUX_DCHECK(initialized_);
        return std::make_shared<DictionaryVector<T> >(
            this->pool_,
            this->slice_nulls(offset, length),
            length,
            value_vector(),
            indices_
                ? Buffer::slice<vector_size_t>(indices_, offset, length, this->pool_)
                : indices_);
    }

    template<typename T>
    void DictionaryVector<T>::validate(const VectorValidateOptions &options) const {
        SimpleVector<T>::validate(options);
        auto indicesByteSize =
                BaseVector::byteSize<vector_size_t>(BaseVector::length_);
        POLLUX_CHECK_GE(indices_->size(), indicesByteSize);
        for (auto i = 0; i < BaseVector::length_; ++i) {
            const bool isNull =
                    BaseVector::rawNulls_ && bits::isBitNull(BaseVector::rawNulls_, i);
            if (isNull) {
                continue;
            }
            // Verify index for a non-null position. It must be >= 0 and < size of the
            // base vector.
            POLLUX_CHECK_GE(
                rawIndices_[i],
                0,
                "Dictionary index must be greater than zero. Index: {}.",
                i);
            POLLUX_CHECK_LT(
                rawIndices_[i],
                dictionaryValues_->size(),
                "Dictionary index must be less than base vector's size. Index: {}.",
                i);
        }
        dictionaryValues_->validate(options);
    }
} // namespace kumo::pollux
